import time
from kafka import KafkaConsumer, KafkaProducer
from kafka import TopicPartition


def count(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        s1 = func(*args, **kwargs)
        time.sleep(1)
        s2 = func(*args, **kwargs)
        end = time.time()
        print(f'增加了{s2 - s1}条记录，耗时{round(end - start)}秒')
        return func

    return wrapper()


def consumer():
    consumer = KafkaConsumer(
        bootstrap_servers=['10.0.17.17:9091', '10.0.17.17:9092', '10.0.17.17:9093'])
    tp = TopicPartition(topic='COMMUNICATION_RESPONSE_DATA', partition=0)
    # tp1 = TopicPartition(topic='COMMUNICATION_REQUEST_DATA', partition=0)
    consumer.assign([tp])
    # start = time.time()
    count = consumer.position(tp)
    consumer.close()
    return count
    # while True:
    #     consumer.assign([tp])
    #     # print(consumer.position(tp))
    #     count2 = consumer.position(tp)
    #     print(f'目前共有{count2}条记录')
    #     end = time.time()
    #     res = end - start
    #     res_count = count2 - count1
    #     print(f'耗时{round(res, 2)}秒,增加{res_count}条记录')
    #     time.sleep(1)


# if __name__ == '__main__':
#     s1 = consumer()
# start = time.time()
# while True:
#     s2 = consumer()
#     time.sleep(1)
#     end = time.time()
#     print(f'耗时{round((end - start), 2)}秒,增加了{s2 - s1}条记录')
# print(s1)
def count_offset():
    con = KafkaConsumer(bootstrap_servers=['10.0.17.17:9091', '10.0.17.17:9092', '10.0.17.17:9093'])
    con.partitions_for_topic('COMMUNICATION_RESPONSE_DATA')
    tps = [TopicPartition('COMMUNICATION_RESPONSE_DATA', p) for p in
           con.partitions_for_topic('COMMUNICATION_RESPONSE_DATA')]
    con.assign(tps)
    sum_total = con.end_offsets(tps)
    # print(sum_total[TopicPartition(topic='COMMUNICATION_RESPONSE_DATA', partition=0)])
    return sum_total[TopicPartition(topic='COMMUNICATION_RESPONSE_DATA', partition=0)]


def send_message_to_kafka():
    producer = KafkaProducer(bootstrap_servers=['10.0.17.17:9091', '10.0.17.17:9092', '10.0.17.17:9093'])
    for i in range(1):
        msg = 'msg is ' + str(i)
        producer.send('COMMUNICATION_RESPONSE_DATA', msg.encode('utf-8'))


if __name__ == '__main__':
    send_message_to_kafka()
    sum_1 = count_offset() - 1
    time_total = 0
    while True:
        time.sleep(1)
        sum_2 = count_offset() - 1
        time_total += 1
        print(f'经过{time_total}秒，目前offset总数为{sum_2}个，增加了offset数量为：{sum_2 - sum_1}')
